from .buffer import AutoBuffer
import cv2
import logging
import numpy as np


class ImageBuffer:
    def __init__(self, file_paths, buffer_size:int=300, logger=None):
        self.file_paths = file_paths
        self.buffer = AutoBuffer(buffer_size, [self.get_next_image])
        self.pos = 0
        self.logger = logging.getLogger() if logger is None else logger

    def get_num_images(self):
        return len(self.file_paths)

    def get_next_image(self):
        if self.pos == self.get_num_images():
            self.pos = 0
        path = self.file_paths[self.pos]
        self.pos += 1
        img = cv2.imread(path)
        if img is None:
            raise Exception('Can not read image', path)
        return img

    def start(self):
        self.buffer.start()

    def stop(self):
        self.buffer.stop()

    def pause(self):
        self.buffer.pause()

    def resume(self):
        self.buffer.resume()

    def next_batch(self, batch_size):
        num = self.get_num_images()
        if batch_size > num:
            batch_size = num
        assert batch_size >= 0, 'batch_size must be great than or equal to 0'
        result = []
        for _ in range(batch_size):
            img = self.buffer.get()
            if img is None:
                break
            result.append(img)
        return result


def get_binary_image(img, threshold, max_val=255, inverse=False):
    if len(img.shape) > 2 and img.shape[2] == 3:
        img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    ret, img = cv2.threshold(img, threshold, max_val, cv2.THRESH_BINARY_INV if inverse else cv2.THRESH_BINARY)
    return img


def get_corners(binary_img, *names):
    """
    返回二值图binary_img中name所指向的边角的坐标
    :param binary_img: 二值图，其中非0表示前景，0表示背景。
    :param names: name的列表，name取值'top-left', 'bottom-left', 'top-right', 'bottom-right', 分别表示图片上的四个边角
    """
    ys, xs = np.nonzero(binary_img)
    points = []
    dx_dy_value = []
    for name in names:
        if 'left' in name:
            dx = 1
        elif 'right' in name:
            dx = -1
        else:
            raise Exception('name("%s")应该含有left或right' % name)
        if 'top' in name:
            dy = 1
        elif 'bottom':
            dy = -1
        else:
            raise Exception('name("%s")应该含有top或bottom' % name)

        found = False
        for old_dx, old_dy, old_value in dx_dy_value:
            if old_dx == -dx and old_dy == -dy:
                i = np.argmax(old_value)
                found = True
                break
        if not found:
            if dx > 0:
                if dy > 0:
                    value = xs + ys
                else:
                    value = xs - ys
            elif dy > 0:
                value = ys - xs
            else:
                value = -(xs + ys)
            dx_dy_value.append((dx, dy, value))
            i = np.argmin(value)
        points.append((xs[i], ys[i]))
    return points


def get_first_point(binary_img, name, foreground=255):
    """
    返回二值图img中name所指向的第一个前景点
    :param binary_img: 二值图
    :param name: 取值'left-top', 'left-bottom', 'right-top', 'right-bottom', 分别表示图片四个角上的点
    :param foreground: 前景色，非0或'white'表示白色，0或'black'表示黑色
    """
    height, width = binary_img.shape
    if 'left' in name:
        x0 = 0
        dx = 1
    elif 'right' in name:
        x0 = width - 1
        dx = -1
    else:
        raise Exception('name("%s")中未找到left或right' % name)
    if 'top' in name:
        y0 = 0
        dy = 1
    elif 'bottom' in name:
        y0 = height - 1
        dy = -1
    else:
        raise Exception('name("%s")中未找到top或bottom' % name)

    if type(foreground) == str:
        foreground = 1 if foreground == 'white' else 0
    for num in range(1, width + height):
        x, y = x0, y0
        for i in range(num):
            if 0 <= x < width and 0 <= y < height and binary_img[y, x] == foreground:
                return x, y
            else:
                x += dx
                y -= dy
        y0 += dy
    return -1, -1

